import os
import sys
from torch.utils.data import DataLoader
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
from data_loader import IrisDataloader


# 初始化神经网络模型
class NN(nn.Module):
    def __init__(self, in_dim, hidden_dim1, hidden_dim2, out_dim):
        """
        初始化神经网络模型
        :param in_dim: 输入维度
        :param hidden_dim1: 第一个隐藏层维度
        :param hidden_dim2: 第二个隐藏层维度
        :param out_dim: 输出维度
        """
        super().__init__()

        # 定义神经网络的层结构(实例化了三个线性层)
        self.layer1 = nn.Linear(in_dim, hidden_dim1)
        self.layer2 = nn.Linear(hidden_dim1, hidden_dim2)
        self.layer3 = nn.Linear(hidden_dim2, out_dim)

    def forward(self, x):
        """
        前向传播函数
        :param x: 输入数据
        :return: 经过神经网络计算后的输出结果
        """

        # 数据通过网络的每一层进行计算
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)

        return x


# 定义计算环境, 如果有可用的 CUDA  GPU，就使用 GPU(第0块)，否则使用 CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 调用自定义数据集加载器, 按照指定路径加载得到自定义数据集对象custom_dataset
custom_dataset = IrisDataloader("./Iris_data.txt")

# 数据集的划分和加载。训练集(70%)，验证集(20%)和测试集(10%)
train_size = int(len(custom_dataset) * 0.7)
val_size = int(len(custom_dataset) * 0.2)
test_size = len(custom_dataset) - train_size - val_size

# 按照指定的数据集大小, 随机划分数据集为训练集、验证集和测试集
train_dataset, val_dataset, test_dataset = (
    torch.utils.data.random_split(custom_dataset, [train_size, val_size, test_size]))

# 构建训练集的数据加载器 DataLoader, 用于批量加载数据集
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)

# 打印数据集的大小
# print("训练集的大小", len(train_loader) * 16, "验证集的大小", len(train_loader), "测试集的大小", len(test_loader))
print("训练集的大小", train_size, "验证集的大小", val_size, "测试集的大小", test_size)


# 推理函数
def infer(model, dataset, device):
    """
    计算模型在给定数据集上的准确率
    :param model: 要进行推理的神经网络模型实例
    :param dataset: 用于推理的数据集
    :param device: 模型和数据要迁移到的计算设备
    :return: 模型在数据集上的准确率
    """
    # 将模型设置为评估模式
    model.eval()
    # 预测正确的样本数量
    acc_num = 0

    # 禁用梯度计算
    with torch.no_grad():
        for data in dataset:
            # 将数据项拆分为输入数据 datas 和对应的标签 label
            datas, label = data
            # 数据传入模型进行前向推理
            outputs = model(datas.to(device))
            # 对模型输出 outputs ，在维度 dim=1 上（一般是类别维度 ）取最大值对应的索引，得到预测的类别索引 predict_y ，这个索引就代表预测的类别
            # output的形状是(batch_size, num_classes)，dim=1表示在类别维度上取最大值
            # torch.max 返回一个元组，第一个元素是最大值，第二个元素是最大值的索引
            predict_y = torch.max(outputs, dim=1)[1]

            # 统计预测正确的样本数量
            # item() 将张量转换为 Python 数值
            acc_num += torch.eq(predict_y, label.to(device)).sum().item()

    # 计算准确率(预测正确的样本数 / 总样本数)
    acc = acc_num / len(dataset)
    return acc


# 主函数，用于执行模型训练等主要流程
def main(lr=0.005, epochs=20):
    """
    主函数，用于执行模型训练等主要流程
    :param lr: 学习率, 控制模型参数更新的步长, 默认值为 0.005
    :param epochs: 训练的轮数默认值为 20
    :return:
    """
    # 实例化神经网络模型 NN
    model = NN(4, 12, 6, 3).to(device)
    # 定义损失函数为交叉熵损失 CrossEntropyLoss
    loss_f = nn.CrossEntropyLoss()

    # 构建参数优化器要优化的参数列表 pg ，筛选出模型中需要计算梯度（requires_grad 为 True ）的参数
    pg = [p for p in model.parameters() if p.requires_grad]
    # 使用 Adam 优化器，传入要优化的参数列表 pg 和学习率 lr ，来对模型参数进行优化
    optimizer = optim.Adam(pg, lr=lr)

    # 权重文件存储路径
    # os.getcwd() 获取当前工作目录
    save_path = os.path.join(os.getcwd(), "results/weights")
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # 开始训练，外层循环，遍历训练的轮数
    for epoch in range(epochs):
        # 将模型设置为训练模式
        model.train()
        # 初始化一个用于统计正确预测数量的张量
        # torch.zeros(1)：形状 [1]（单元素）。
        acc_num = torch.zeros(1).to(device)
        # 用于统计当前轮次训练时处理的样本总数，初始化为 0
        sample_num = 0

        # 使用 tqdm 库创建一个进度条 train_bar ，用于显示训练过程的进度
        train_bar = tqdm(train_loader, file=sys.stdout, ncols=100)
        # 内层循环，遍历训练数据加载器中的每个批次数据
        for datas in train_bar:
            # 将每个批次的数据 datas 拆分为输入数据 data 和对应的标签 label
            data, label = datas
            # 挤压标签的最后一个维度，如果标签维度是 (batch_size, 1) 等情况，挤压后变成 (batch_size) ，方便后续计算
            # 表示从最后一个维度开始检查并移除大小为 1 的维度。
            label = label.squeeze(-1)
            # 累加当前批次的样本数量，data.shape[0] 就是当前批次的样本数（batch_size ）
            sample_num += data.shape[0]

            # 优化器梯度清零，因为 PyTorch 中梯度会累加，所以每次训练前要清除之前的梯度
            optimizer.zero_grad()
            # 将输入数据 data 迁移到指定设备（device）上，然后传入模型进行前向传播，得到输出 outputs
            outputs = model(data.to(device))
            # 对模型输出 outputs ，在维度 dim=1 上（类别维度 ）取最大值对应的索引，得到预测的类别 pred_class
            # torch.max 返回一个元组，第一个元素是最大值，第二个元素是最大值的索引，这里取索引就是预测类别
            pred_class = torch.max(outputs, dim=1)[1]
            # 将预测类别 pred_class 和标签 label（迁移到相同设备 ）进行比较，统计相等（预测正确 ）的数量，累加到 acc_num 中
            acc_num += torch.eq(pred_class, label.to(device)).sum()

            # 计算损失，将模型输出 outputs 和标签 label（迁移到设备 ）传入损失函数 loss_f 中
            loss = loss_f(outputs, label.to(device))
            # 反向传播，计算损失关于模型参数的梯度
            loss.backward()
            # 优化器执行一步参数更新，根据计算得到的梯度来更新模型的参数
            optimizer.step()

            # 训练的准确率=正确预测数量/已处理样本总数
            train_acc = acc_num / sample_num
            # 设置进度条的描述信息，显示当前是第几轮训练（epoch + 1 是因为 epoch 从 0 开始计数的话，这样显示更直观 ）、总轮数、当前损失值
            train_bar.desc = "train epoch[{}/{}] loss:{:.3f} ".format(epoch + 1, epochs, loss)

        # 调用 infer 函数计算验证集上的准确率，传入模型、验证集数据加载器 val_loader 和设备 device
        val_acc = infer(model, val_loader, device)
        # 打印当前轮次的训练信息，包括轮次、总轮数、损失值、训练集准确率、验证集准确率
        print(
            "train epoch:{}/{} loss:{:.3f} train_acc:{:.3f} val_acc:{:.3f}".
            format(epoch + 1, epochs, float(loss), float(train_acc), float(val_acc)))
        # 保存当前模型的状态字典（即模型参数 ）到指定路径 save_path 下的 "nn.pth" 文件中
        torch.save(model.state_dict(), os.path.join(save_path, "nn.pth"))

        # 每次数据集迭代（一轮训练 ）之后，要对初始化的指标清零，为下一轮训练做准备
        train_acc = 0.
        val_acc = 0.
        acc_num = torch.zeros(1).to(device)

    print("Finished Training")

    # 调用 infer 函数计算测试集上的准确率，传入模型、测试集数据加载器 test_loader 和设备 device
    test_acc = infer(model, test_loader, device)
    print("test_acc:", test_acc)


if __name__ == '__main__':
    main()
